Implement bin sort with minsigma - #196

use std::hint::black_box;
use std::time::{Duration, Instant};

use rustuna_core::storage::InMemoryStorage;
use rustuna_core::study::{create_study, Direction};
use rustuna_core::Result;
use rustuna_sampler::tpe::TpeSampler;

const N_TRIALS: [usize; 3] = [500, 1000, 2000];
const N_PARAMS: [usize; 1] = [40];

fn run_study(n_trials: usize, n_params: usize) -> Result<Duration> {
    let storage = InMemoryStorage::new();
    let study = create_study(
        "tpe-benchmark",
        storage,
        TpeSampler::seed_from_u64(0),
        vec![Direction::Minimize],
    )?;

    let start = Instant::now();
    study.optimize(
        |mut trial| {
            let mut value = 0.0;
            for i in 0..n_params {
                let x = trial.suggest_float(&format!("x{i}"), -10.0, 10.0)?;
                value += x * x;
            }
            Ok(vec![black_box(value)])
        },
        n_trials,
    )?;
    Ok(start.elapsed())
}

fn main() -> Result<()> {
    println!(
        "{:>8}  {:>8}  {:>12}  {:>12}",
        "n_params", "n_trials", "total [ms]", "per trial [us]"
    );
    for n_params in N_PARAMS {
        for n_trials in N_TRIALS {
            let duration = run_study(n_trials, n_params)?;
            println!(
                "{:>8}  {:>8}  {:>12.1}  {:>12.1}",
                n_params,
                n_trials,
                duration.as_secs_f64() * 1e3,
                duration.as_secs_f64() * 1e6 / n_trials as f64,
            );
        }
    }
    Ok(())
}